//
// Created by wchenbt on 5/4/2021.
//
#ifdef ENABLE_MAPLE_SAN

#include "san_common.h"

#include "asan_interfaces.h"
#include "me_function.h"
#include "me_ir.h"
#include "mir_builder.h"
#include "string_utils.h"

namespace maple {

std::vector<StmtBlockNodePair> getAllOrderedStmtNodeList(BlockNode* block) {
  // When optimization enabled, there are BlockNode with list of StmtNode
  // We recover the recursive structure into linear structure
  // We could later optimize the recursive implementation with a stack
  std::vector<StmtBlockNodePair> ret;
  for (StmtNode &stmt : block->GetStmtNodes()) {
    if (stmt.GetOpCode() == OP_block) {
      BlockNode* tmpBlock = dynamic_cast<BlockNode*>(&stmt);
      CHECK_FATAL((tmpBlock != nullptr), "Fail to convert stmt to BlockNode");
      std::vector<StmtBlockNodePair> tmp = getAllOrderedStmtNodeList(tmpBlock);
      ret.reserve(ret.size() + tmp.size());
      ret.insert(ret.end(), tmp.begin(), tmp.end());
    }
    else {
      ret.push_back(StmtBlockNodePair(&stmt, block));
      for (size_t i = 0; i < stmt.NumOpnds(); ++i) {
        StmtNode* tmpStmt = dynamic_cast<StmtNode*>(stmt.Opnd(i));
        if (tmpStmt == nullptr) continue;
        if (tmpStmt->GetOpCode() == OP_block) {
          BlockNode* tmpBlock = dynamic_cast<BlockNode*>(tmpStmt);
          CHECK_FATAL((tmpBlock != nullptr), "Fail to convert stmt to BlockNode");
          std::vector<StmtBlockNodePair> tmp = getAllOrderedStmtNodeList(tmpBlock);
          ret.reserve(ret.size() + tmp.size());
          ret.insert(ret.end(), tmp.begin(), tmp.end());
        }
      }
    }
  }
  return ret;
}

void appendToGlobalCtors(const MIRModule &mirModule, const MIRFunction *func) {
  MIRBuilder *mirBuilder = mirModule.GetMIRBuilder();
  MIRFunction *GlobalCtors = mirBuilder->GetOrCreateFunction("__cxx_global_var_init", TyIdx(PTY_void));
  MapleVector<BaseNode *> args(mirBuilder->GetCurrentFuncCodeMpAllocator()->Adapter());
  CallNode *callNode = mirBuilder->CreateStmtCall(func->GetPuidx(), args);
  GlobalCtors->GetBody()->AddStatement(callNode);
}

void appendToGlobalDtors(const MIRModule &mirModule, const MIRFunction *func) {
  MIRBuilder *mirBuilder = mirModule.GetMIRBuilder();
  MIRFunction *GlobalCtors = mirBuilder->GetOrCreateFunction("__cxx_global_var_fini", TyIdx(PTY_void));
  MapleVector<BaseNode *> args(mirBuilder->GetCurrentFuncCodeMpAllocator()->Adapter());
  CallNode *callNode = mirBuilder->CreateStmtCall(func->GetPuidx(), args);
  GlobalCtors->GetBody()->AddStatement(callNode);
}

MIRFunction *getOrInsertFunction(MIRBuilder *mirBuilder, const char *name, MIRType *retType,
                                 std::vector<MIRType *> argTypes) {
  GStrIdx strIdx = GlobalTables::GetStrTable().GetStrIdxFromName(name);
  MIRFunction *func = mirBuilder->GetOrCreateFunction(name, retType->GetTypeIndex());

  if (strIdx != 0u) {
    return func;
  }

  func->AllocSymTab();

  /* use void* for PTY_dynany */
  if (retType->GetPrimType() == PTY_dynany) {
    retType = GlobalTables::GetTypeTable().GetPtr();
  }

  std::vector<MIRSymbol *> formals;
  for (uint32 j = 0; j < argTypes.size(); ++j) {
    MIRType *argTy = argTypes.at(j);
    /* use void* for PTY_dynany */
    if (argTy->GetPrimType() == PTY_dynany) {
      argTy = GlobalTables::GetTypeTable().GetPtr();
    }
    MIRSymbol *argSt = func->GetSymTab()->CreateSymbol(kScopeLocal);
    const uint32 bufSize = 18;
    char buf[bufSize] = {'\0'};
    int eNum = sprintf_s(buf, bufSize - 1, "p%u", j);
    if (eNum == -1) {
      FATAL(kLncFatal, "sprintf_s failed");
    }
    std::string strBuf(buf);
    argSt->SetNameStrIdx(mirBuilder->GetOrCreateStringIndex(strBuf));
    argSt->SetTyIdx(argTy->GetTypeIndex());
    argSt->SetStorageClass(kScFormal);
    argSt->SetSKind(kStVar);
    func->GetSymTab()->AddToStringSymbolMap(*argSt);
    formals.emplace_back(argSt);
  }
  func->SetAttr(FuncAttrKind::FUNCATTR_public);
  func->SetAttr(FuncAttrKind::FUNCATTR_extern);
  func->UpdateFuncTypeAndFormalsAndReturnType(formals, retType->GetTypeIndex(), false);
  return func;
}

/// Create a global describing a source location.
MIRAddrofConst *createSourceLocConst(MIRModule &mirModule, MIRSymbol *Var, PrimType primType) {
  MIRStrConst *moduleName = createStringConst(mirModule, mirModule.GetFileName(), PTY_a64);
  MIRConst *LocData[] = {
      moduleName,
      GlobalTables::GetIntConstTable().GetOrCreateIntConst(Var->GetSrcPosition().LineNum(),
                                                           *GlobalTables::GetTypeTable().GetInt32()),
      GlobalTables::GetIntConstTable().GetOrCreateIntConst(Var->GetSrcPosition().Column(),
                                                           *GlobalTables::GetTypeTable().GetInt32()),
  };
  // Create struct type
  MIRStructType LocStruct(kTypeStruct);
  GlobalTables::GetTypeTable().AddFieldToStructType(LocStruct, "module_name",
                                                    *GlobalTables::GetTypeTable().GetTypeFromTyIdx(TyIdx(primType)));
  GlobalTables::GetTypeTable().AddFieldToStructType(LocStruct, "line", *GlobalTables::GetTypeTable().GetInt32());
  GlobalTables::GetTypeTable().AddFieldToStructType(LocStruct, "column", *GlobalTables::GetTypeTable().GetInt32());
  // Create initial value
  MIRAggConst *LocStructConst = mirModule.GetMemPool()->New<MIRAggConst>(mirModule, LocStruct);
  // Initialize the field orig
  for (uint32_t i = 0; i < 3; i++) {
    LocStructConst->AddItem(LocData[i], i + 1);
  }
  // Create a new symbol, MIRStructType is a subclass of MIRType
  TyIdx LocStructTy = GlobalTables::GetTypeTable().GetOrCreateMIRType(&LocStruct);
  MIRSymbol *LocStructSym = mirModule.GetMIRBuilder()->CreateSymbol(LocStructTy, Var->GetName() + "_Loc", kStConst,
                                                                    kScGlobal, nullptr, kScopeGlobal);
  LocStructSym->SetKonst(LocStructConst);
  return createAddrofConst(mirModule, LocStructSym, primType);
}

MIRAddrofConst *createAddrofConst(const MIRModule &mirModule, const MIRSymbol *mirSymbol, PrimType primType) {
  AddrofNode *addrofNode = mirModule.GetMIRBuilder()->CreateAddrof(*mirSymbol);
  MIRAddrofConst *mirAddrofConst =
      mirModule.GetMemPool()->New<MIRAddrofConst>(addrofNode->GetStIdx(), addrofNode->GetFieldID(),
                                                  *GlobalTables::GetTypeTable().GetTypeFromTyIdx(TyIdx(primType)));
  return mirAddrofConst;
}

// Create a constant for Str so that we can pass it to the run-time lib.
MIRStrConst *createStringConst(const MIRModule &mirModule, const std::basic_string<char>& Str, PrimType primType) {
  MIRStrConst *strConst =
      mirModule.GetMemPool()->New<MIRStrConst>(GlobalTables::GetUStrTable().GetOrCreateStrIdxFromName(Str),
                                               *GlobalTables::GetTypeTable().GetTypeFromTyIdx(TyIdx(primType)));

  return strConst;
}

bool isTypeSized(MIRType *type) {
  if (type->GetKind() == kTypeScalar || type->GetKind() == kTypePointer || type->GetKind() == kTypeBitField) {
    return true;
  }
  if (type->GetKind() != kTypeStruct && type->GetKind() != kTypeStructIncomplete && type->GetKind() != kTypeArray &&
      type->GetKind() != kTypeFArray && type->GetKind() != kTypeUnion) {
    return false;
  }
  if (type->GetKind() == kTypeArray) {
    MIRArrayType *arrayType = dynamic_cast<MIRArrayType *>(type);
    if (arrayType) {
      return isTypeSized(arrayType->GetElemType());
    }
  }
  if (type->GetKind() == kTypeFArray) {
    MIRFarrayType *farrayType = dynamic_cast<MIRFarrayType *>(type);
    if (farrayType) {
      return isTypeSized(farrayType->GetElemType());
    }
  }
  if (type->IsStructType()) {
    MIRStructType *structType = dynamic_cast<MIRStructType *>(type);
    if (structType) {
      for (size_t i = 1; i < structType->GetFieldsSize(); i++) {
        // FieldID type is int32, there should be a FieldID_MAX
        CHECK_FATAL(i < INT32_MAX, "Too large filed size.");
        if (!isTypeSized(structType->GetFieldType(FieldID(i)))) {
          return false;
        }
      }
      return true;
    }
  }
  return false;
}

std::vector<MIRSymbol *> GetGlobalVaribles(const MIRModule &mirModule) {
  std::vector<MIRSymbol *> globalVarVec;
  for (auto sit = mirModule.GetSymbolDefOrder().begin(); sit != mirModule.GetSymbolDefOrder().end(); ++sit) {
    MIRSymbol *s = GlobalTables::GetGsymTable().GetSymbolFromStidx((*sit).Idx());
    CHECK_FATAL(s != nullptr, "nullptr check");
    if (s->IsJavaClassInterface()) {
      continue;
    }
    if (!s->IsDeleted() && !s->GetIsImported() && !s->GetIsImportedDecl()) {
      if (s->GetSKind() == kStVar) {
        globalVarVec.push_back(s);
      }
    }
  }
  return globalVarVec;
}

int computeRedZoneField(MIRType *type) {
  int field = 1;
  if (type->IsStructType()) {
    MIRStructType *structType = dynamic_cast<MIRStructType *>(type);
    for (size_t i = 1; i < structType->GetFieldsSize(); i++) {
      CHECK_FATAL(i < INT32_MAX, "Too large field size.");
      MIRType *subType = structType->GetFieldType(FieldID(i));
      field += computeRedZoneField(subType);
    }
  }
  return field;
}

size_t TypeSizeToSizeIndex(uint32_t TypeSize) {
  uint32_t Val = TypeSize / 8;
  if (!Val) {
    return std::numeric_limits<uint32_t>::digits;
  }
  if (Val & 0x1) {
    return 0;
  }

  // Bisection method.
  unsigned zeroBits = 0;
  uint32_t shift = std::numeric_limits<uint32_t>::digits >> 1;
  uint32_t mask = std::numeric_limits<uint32_t>::max() >> shift;
  while (shift) {
    if ((Val & mask) == 0) {
      Val >>= shift;
      zeroBits |= shift;
    }
    shift >>= 1;
    mask >>= shift;
  }
  assert(zeroBits < kNumberOfAccessSizes);
  return zeroBits;
}

MIRSymbol *getOrCreateSymbol(MIRBuilder *mirBuilder, const TyIdx tyIdx, const std::string &name, MIRSymKind mClass,
                             MIRStorageClass sClass, MIRFunction *func, uint8 scpID) {
  MIRSymbol *st = nullptr;
  if (func) {
    st = func->GetSymTab()->GetSymbolFromStrIdx(mirBuilder->GetOrCreateStringIndex(name));
  } else {
    st = GlobalTables::GetGsymTable().GetSymbolFromStrIdx(mirBuilder->GetOrCreateStringIndex(name));
  }
  if (st == nullptr || st->GetTyIdx() != tyIdx) {
    return mirBuilder->CreateSymbol(tyIdx, name, mClass, sClass, func, scpID);
  }
  CHECK_FATAL(mClass == st->GetSKind(),
             "trying to create a new symbol that has the same name and GtyIdx. might cause problem");
  CHECK_FATAL(sClass == st->GetStorageClass(),
              "trying to create a new symbol that has the same name and tyIdx. might cause problem");
  return st;
}

CallNode* CreateStackCheck(MIRBuilder *mirBuilder, BaseNode *baseNode, BaseNode *sizeNode) {
  MIRType *voidType = GlobalTables::GetTypeTable().GetVoid();
  MIRType *intType = GlobalTables::GetTypeTable().GetUInt64();
  std::vector<MIRType*> argTypes = {intType, intType};
  // Initialize the function
  MIRFunction * check_func = getOrInsertFunction(mirBuilder, kAsanStackCheck2, voidType, argTypes);
  check_func->SetAttr(FUNCATTR_extern);
  MapleVector<BaseNode*> args(mirBuilder->GetCurrentFuncCodeMpAllocator()->Adapter());
  args.push_back(baseNode);
  args.push_back(sizeNode);
  return mirBuilder->CreateStmtCall(kAsanStackCheck2, args);
}

// Code for Sanrazor
int SANRAZOR_MODE() {
  /*
  Sanrazor has several mode
  0. Didn't turn on / default
  1. intrument for coverage
  2. collected coverage, analysis and remove sanitzer check 
  */
  char *env = getenv("SANRAZOR_MODE");
  int SanrazorMode = 0;
  if (env) {
    SanrazorMode = atoi(env);
    if (SanrazorMode >= 3) {
      return 0;
    }
    return SanrazorMode;
  } else {
    return 0;
  }
}

}  // namespace maple

#endif